# cython: profile=True
from numpy import dot,zeros,array,ones,argmax,zeros_like,shape,isnan
#from math import sqrt,sin,cos,exp,pi
import cython,time, pdb
if not cython.compiled:
    from math import *
from v3dsolver.fluid_elements.fluid_elems import Saffman
from v3dsolver.neighbour_finder import Bin,NeighbourFinder
from v3dsolver.vector cimport Vec,Vec_new
from lpsolver cimport csolver
from random import choice
from mpi4py import MPI
from cython.operator cimport dereference as deref

cdef extern from "stdlib.h":
    void free(void* ptr)
    void* malloc(size_t size)
    void* realloc(void* ptr,size_t size)
    
cdef unitsphere_possible_nbrs(int n):
    """
    Returns a list of uniformly spaced points on a unit sphere centered 
    at the origin. 
        
    Any new neighbour would be placed at one of these lat-long positions.
        
    n: number of new points on each latitude (same as no. of such latitudes)
    """
    cdef list pts = [] 
    cdef int i,j
    cdef double x,y,z
    cdef double PI = 3.1415926535897931
    cdef double dtheta = 2.*PI/n
    cdef double theta,phi
    cdef Vec newpt
    for i in range(n/2+1):
        theta=dtheta*i
        z = cos(theta)
        if i==0 or i==(n-1):
            # just add one particle at the pole
            x = 0.0
            y = 0.0
            newpt = Vec_new(x,y,z)
            pts.append(newpt)
        else:
            for j in range(n):
                phi=dtheta*j
                x=sin(theta)*cos(phi)
                y=sin(theta)*sin(phi)
                newpt = Vec_new(x,y,z)
                pts.append(newpt)
    return pts

cdef list kernel_nbrs = unitsphere_possible_nbrs(8)

cdef bint check_intersection(Vec pt1,Vec pt2,panel):
    """
    Returns true if line joining pt1 and pt2 intersects panel.
    """
    return False # to be implemented

cdef class VRTSolver:
    """
    This class is responsible for solving VRT around an element and providing
    vorticities obtained by the neighbours of that element.
    """
    def __init__(self,base_elem,neighbours,finebin,cellsize,n=8):
        """
        Initialization of VRTSolver class. 
        
        base_elem: element for which VRT is to be solved (of form: [hash,elem])
        neighbours: list of neighbours to base_elem (each of form: [hash,elem])
        finebin: a bin of size half the original bin containing all elems
                 (including preadded elems)
        cellsize: criteria for neighbourhood: 
             sqrt(8*nu*dt): for 3D 
             sqrt(12*nu*dt): for 2D
        n: number of new possible neighbours on each latitude (same as no. of 
          such latitudes)
        """
        self.base_elem = base_elem
        self.base_index = int(self.base_elem[0].split()[1]) # base_elem index
        self.nbrs = neighbours
        self.cellsize = cellsize
        self.h_nu = self.cellsize/sqrt(8) ## only for 3D
        self.n = n
        self.possible_nbrs = self.get_possible_nbrs(n) 
        self.finebin = finebin
        self.len_elems = len(self.finebin.elems)
        self._hash_prefix = str(MPI.COMM_WORLD.rank) + ' '
    
    def get_possible_nbrs(self,n):
        """
        Returns a list of uniformly spaced points on a sphere centered at base
        elem and of radius cellsize. 
        
        Any new neighbour would be placed at one of these lat-long positions, 
        depending on distance of this new neighbour from other neighbours.
        
        n: number of new points on each latitude (same as no. of such latitudes)
        """
        global kernel_nbrs
        cdef list pts = []
        cdef Vec base_elem_pos = self.base_elem[1].pos
        cdef Vec kernel_nbr 
        cdef Vec new_nbr
        for i in range(len(kernel_nbrs)):
            kernel_nbr = kernel_nbrs[i]
            new_nbr = kernel_nbr*self.cellsize + base_elem_pos
            pts.append(new_nbr)
        return pts
    
    def check_bdry_intersection(self,nbrpos,nbrcell):
        """
        Checks if nbrpos is an acceptable candidate for new nbr by 
        checking if it lies inside the boundary.
        """
        basepos = self.base_elem[1].pos
        basecell = self.finebin.get_index(basepos)
        test_panels = (self.finebin.panels[basecell] + 
                       self.finebin.panels[nbrcell])
        for panel in test_panels:
            if check_intersection(basepos,nbrpos,panel):
                return True
        return False
        
    
    cdef add_neighbour(self):
        """
        Adds a neighbour particle to base_elem in case the number of neighbours
        it already has is less than the number of equations of VRT. 
        
        It tries to add the neighbour such that the sum of its distance from 
        all the other neighbours is maximum. This approach tends to fill up 
        holes in the vorticity distribution.
        """
        # sum of distance of possible neighbours from all other neighbours
        l = len(self.possible_nbrs)
        dist_possible_nbrs = zeros(l,dtype='float64')
        
        for j in range(len(self.nbrs)):
            nbr = self.nbrs[j]
            nbr_pos = nbr[1]
            for i in range(l):
                possible_nbr = self.possible_nbrs[i]
                rvec = possible_nbr - nbr_pos
                r = rvec.length()
                dist_possible_nbrs[i] += r 
        argnbr = argmax(dist_possible_nbrs)
        #argnbr = choice(range(l))
        new_pos = self.possible_nbrs[argnbr]
        
        """
        parentcell = self.finebin.get_index(self.base_elem[1].pos)
        index_1d = (parentcell.z*self.finebin.ncells_x*self.finebin.ncells_y + 
                    parentcell.y*self.finebin.ncells_x + parentcell.x)
        least_population = self.finebin.elem_count[index_1d]
        candidate_nbrcell = index_1d
        for nbr in self.finebin.nbrlist[index_1d]:
            if not self.finebin.elem_count.has_key(nbr):
                candidate_nbrcell = nbr
                break
            if self.finebin.elem_count[nbr] < least_population:
                least_population = self.finebin.elem_count[nbr]
                candidate_nbrcell = nbr
        new_pos = self.finebin.get_newelem_pos(candidate_nbrcell)
        """
        
        new_index = self.len_elems + self.base_index
        self.base_index = new_index
        str_new_index = str(new_index)
        new_hash = self._hash_prefix + str_new_index
        self.nbrs.append([new_hash, new_pos])

    cdef get_fractions(self,double* f,double* pmax_f,double* pmin_f,double* pfsum):
        """
        Returns the fractions of vorticity obtained by solving the Vorticity 
        Redistribution Technique equations.
        """
        cdef int neq = 10              # number of VRT equations
        cdef int nvar = len(self.nbrs) # number of variables = number of neighbours
        cdef int i,j

        cdef double** Aeq = <double**>malloc(neq*sizeof(double))
        for i in range(neq):
            Aeq[i] = <double*>malloc(nvar*sizeof(double))
        cdef double* beq = <double*>malloc(neq*sizeof(double))
        
        cdef Vec eps
        cdef double eps1
        cdef double eps2
        cdef double eps3
        cdef Vec base_elem_pos = self.base_elem[1].pos
        cdef double h_nu = self.h_nu
        cdef double R = sqrt(8)
        cdef Vec nbr_pos

        # making equation in terms of f_ij
        for j in range(nvar):
            nbr_pos = self.nbrs[j][1]
            eps=(base_elem_pos - nbr_pos)/h_nu
            eps2 = eps.y
            eps1 = eps.x
            eps3 = eps.z
            Aeq[0][j] = 1.
            Aeq[1][j] = eps1
            Aeq[2][j] = eps2
            Aeq[3][j] = eps3
            Aeq[4][j] = eps1*eps2
            Aeq[5][j] = eps2*eps3
            Aeq[6][j] = eps1*eps3
            Aeq[7][j] = eps1*eps1
            Aeq[8][j] = eps2*eps2
            Aeq[9][j] = eps3*eps3
        
        beq[0] = 1.  
        beq[1] = 0.
        beq[2] = 0.
        beq[3] = 0.
        beq[4] = 0.
        beq[5] = 0.
        beq[6] = 0.   
        beq[7] = 2.
        beq[8] = 2.
        beq[9] = 2.
        
        cdef double sum_Aeq_i
        for i in range(neq):
            sum_Aeq_i = 0.
            for j in range(nvar):
                sum_Aeq_i += Aeq[i][j]
            beq[i] = beq[i] - 0.5*sum_Aeq_i
        
        cdef double* w = <double*>malloc((nvar+1)*sizeof(double))
        csolver(Aeq,beq,w,neq,nvar)
        
        for j in range(nvar+1):
            f[j] = w[j] + 0.5
    
        cdef double max_f = deref(pmax_f)
        cdef double min_f = deref(pmin_f)
        cdef double fsum = deref(pfsum)
        pmax_f[0] = 0.
        pmin_f[0] = 1.
        pfsum[0] = 0.
        for j in range(nvar):
            if f[j]>pmax_f[0]:
                pmax_f[0] = f[j]
            if f[j]<pmin_f[0]:
                pmin_f[0] = f[j]
            pfsum[0] += f[j]
        f[nvar] = pmax_f[0]
        
        for i in range(neq):
            free(Aeq[i])
        free(Aeq)
        free(beq)
        free(w)
    
    def solve_vrt(self):
        """
        Solves the VRT and returns a list of neighbours in the form:
        [nbrhash,nbrpos,nbrvor]; where nbrvor is component of vorticity 
        obtained by diffusing current base_elem.
        """
        cdef int nvar = len(self.nbrs)
        cdef double min_f
        cdef double max_f
        cdef double fsum
        
        countnew = 0
        countfailed = 0
        
        while nvar<11:
            countnew += 1
            self.add_neighbour()
            nvar += 1
        
        cdef double* f = <double*>malloc((nvar+1)*sizeof(double))
        self.get_fractions(f,&max_f,&min_f,&fsum)
        
        while max_f>1 or min_f<0 or abs(fsum-1.0)>1e-4:
            free(f)
            countfailed += 1
            self.add_neighbour()
            nvar += 1
            f = <double*>malloc((nvar+1)*sizeof(double))
            self.get_fractions(f,&max_f,&min_f,&fsum)
        result = []
        cdef list nbr
        cdef Vec base_elem_vor = self.base_elem[1].vor
        cdef str nbrhash
        cdef Vec nbrpos
        cdef Vec nbrvor_new
        cdef int i
        for i in range(nvar):
            nbr = self.nbrs[i]
            nbrhash = nbr[0]
            nbrpos = nbr[1]
            nbrvor_new = f[i]*base_elem_vor
            result.append([nbrhash,nbrpos,nbrvor_new])
        free(f)
        return result,countnew,countfailed

class DiffusionManager:
    """
    This class is responsible for diffusion of vorticity of all elements. It 
    uses Vorticity Redistribution Technique (VRT) to simulate diffusion.
    """
    def __init__(self,elems,nu,dt,R=sqrt(8)):
        """
        Initialization of DiffusionManager object.
        
        elems: list of elements to be diffused (only local elems in case of 
               parallel implementation)
        nu: kinematic viscosity of the fluid
        dt: timestep for simulation
        R: radius for binning (sqrt(8) for 3D flow; sqrt(12) for 2D flow)
        """
        self.elems = elems
        self.nu = nu
        self.dt = dt
        self.R = R
        self.cellsize = self.R*sqrt(nu*dt)
        
    def get_preadded_elems(self,elems):
        """
        Creates a finer bin based on given elems which is used to add
        neighbours in that part of the bin which is relatively empty. 
        Returns the list of new elems added in this way.
        """
        preadded_positions = []
        cdef Bin finebin = Bin(elems,self.cellsize/2.)
        is_done = {}
        cutoff = 0.5
        for i in finebin.first_elem_indices.keys():
            for j in finebin.nbrlist[i]:
                if not is_done.has_key(j):
                    is_done[j] = 1 # mark cell j as done
                    nbrelems_count = 0
                    if finebin.nbrlist.has_key(j):
                        nbrlist = finebin.nbrlist[j]
                    else:
                        nbrlist = finebin.get_nbrs(j)
                    for k in nbrlist:
                        try:
                            nbrelems_count += finebin.elem_count[k]
                        except KeyError:
                            pass
                    expected_elemcount = nbrelems_count/27.
                    #print 'expected_elemcount:', expected_elemcount
                    try:
                        imbalance = expected_elemcount - finebin.elem_count[j]
                    except KeyError:
                        imbalance = expected_elemcount
                    n_new = int(round(imbalance/cutoff))
                    #print 'n_new:', n_new
                    if n_new>0:
                        for l in range(n_new):
                            # randomly add a particle
                            newpos = finebin.get_newelem_pos(j)
                            preadded_positions.append(newpos)
        
        blob_rad = elems[0].rad
        preadded_elems = []
        for i in range(len(preadded_positions)):
            newelem = Saffman(pos=preadded_positions[i],vor=Vec(),
                              vel=Vec(),rad=blob_rad)
            preadded_elems.append(newelem)
        return preadded_elems

    def dispose_nbrs(self,new_nbrs,comm=MPI.COMM_WORLD):
        """
        Dispose obtained neighbours to appropriate locations. 
        
        * Remote processor vorticities are sent back
        * Vorticities obtained from remote processor are updated 
        * Already existing elems update their vorticities according to local 
          diffusion
        * Newly created elems are added to the same processor
        """
        existing_elems = self.elems + self.preadded_elems
        blob_rad = existing_elems[0].rad
        
        send_data = []
        for i in range(comm.size):
            send_data += [[]]
        
        newelems = []
        newvors = [Vec() for elem in existing_elems]
        
        for new_nbr in new_nbrs:
            new_hash,new_pos,new_vor = new_nbr
            
            # if belonging to remote processor
            if int(new_hash.split()[0])!=comm.rank:
                send_data[int(new_hash.split()[0])] += [[new_hash,new_vor]]
            
            # if newly created
            elif int(new_hash.split()[1])>=len(existing_elems):
                newelems += [Saffman(pos=new_pos,vel=Vec(),
                             vor=new_vor,rad=blob_rad)]
            
            # if already existing
            else:
                index = int(new_hash.split()[1])
                newvors[index] += new_vor
        
        # send non-local data
        reqlist = []
        for i in range(len(send_data)):
            if not i==comm.rank:
                reqlist += [comm.isend(send_data[i],dest=i,tag=7)]
        
        # recv non-local data
        for i in range(comm.size):
            if not i==comm.rank:
                recvd_data = comm.recv(source=i,tag=7)
                for remote_nbr in recvd_data:
                    hash,vor = remote_nbr
                    assert int(hash.split()[0])==comm.rank
                    index = int(hash.split()[1])
                    newvors[index] += vor
        MPI.Prequest.Waitall(reqlist)
        # updating elem vors based on obtained (local and non-local) vors
        for i,elem in enumerate(existing_elems):
            elem.vor = newvors[i]
        
        # adding newly created elems to existing elems list
        self.elems += self.preadded_elems + newelems

    def diffuse(self,comm=MPI.COMM_WORLD):
        """
        Perform diffusion using VRT. Returns the list of elems obtained after
        diffusion.
        """
        self.preadded_elems = self.get_preadded_elems(self.elems)
        print 'len(preadded_elems):', len(self.preadded_elems)
        conditioned_elemslist = self.elems+self.preadded_elems
        self.bin = Bin(conditioned_elemslist,self.cellsize)
        self.finebin = Bin(conditioned_elemslist,self.cellsize/2.)
        self.nbrfinder = NeighbourFinder(self.bin,self.nu,self.dt,self.R)
        new_nbrs = []
        len_elems = len(conditioned_elemslist)
        total_countnew, total_countfailed = 0,0
        for i,elem in enumerate(self.elems):
            nbrs = self.nbrfinder.get_neighbours_parallel(elem,comm)
            #print i, len(nbrs),
            base_elem = [self.bin.hashdict[elem],elem]
            vrtsolver = VRTSolver(base_elem,nbrs,self.finebin,self.cellsize)
            vrt_nbrs,countnew,countfailed = vrtsolver.solve_vrt()
            #print countnew, countfailed
            total_countnew += countnew
            total_countfailed += countfailed
            new_nbrs += vrt_nbrs
        print 'total_countnew:', total_countnew
        print 'total_countfailed:', total_countfailed
        print 'total_new:', total_countnew+total_countfailed
        self.dispose_nbrs(new_nbrs,comm)
